const fs = require('fs').promises;
const path = require('path');
const BaseDatasetWrapper = require('../../BaseDatasetWrapper');
const { DatasetListDB } = require('../../../dbs/dbDatasetList');
const { getAllFilesInDirectory } = require('../../../utils/fileUtils');

class PaddleDatasetWrapper extends BaseDatasetWrapper {
    async parse(dataset_info) {
        const datasetId = dataset_info.dataset_id;
        const datasetDir = dataset_info.dataset_dir;

        // 获取所有图片文件和标注文件
        const imageFiles = await getAllFilesInDirectory(datasetDir, ['.jpg', '.jpeg', '.png', '.bmp'], ['workspace']);
        let labelFiles = await getAllFilesInDirectory(datasetDir, ['.txt'], ['workspace']);
        const selections = JSON.parse(dataset_info.label_selections);
        if (selections.length > 0) {
            labelFiles = selections.map((file) => {
                return path.join(datasetDir, file);
            });
        }

        // 创建图像路径映射
        const imagePathMap = new Map();
        imageFiles.forEach(file => {
            const relativePath = path.relative(datasetDir, file);
            imagePathMap.set(relativePath, file);
        });
        // console.log(`imagePathMap: `, imagePathMap);

        // 创建工作目录
        const workDir = path.join(datasetDir, 'workspace', datasetId);
        await fs.mkdir(workDir, { recursive: true });

        // 解析每个标注文件
        const matchingPairs = [];
        let labeledCount = 0;
        for (const labelFile of labelFiles) {
            const content = await fs.readFile(labelFile, 'utf8');
            const lines = content.split('\n').filter(line => line.trim());
            const labelDir = path.dirname(labelFile);

            for (const line of lines) {
                let [imgRelativePath, jsonStr] = line.split('\t');
                if (!imgRelativePath || !jsonStr) continue;

                try {
                    const fullImagePath = path.join(labelDir, imgRelativePath);
                    const datasetRelativePath = path.relative(datasetDir, fullImagePath);
                    const imagePath = imagePathMap.get(datasetRelativePath);
                    // console.log(`fullImagePath: ${fullImagePath}`);
                    // console.log(`datasetRelativePath: ${datasetRelativePath}`);
                    // console.log(`imagePath: ${imagePath}`);

                    if (!imagePath) {
                        console.warn(`Image not found: ${fullImagePath}`);
                        continue;
                    }

                    // 解析标注JSON
                    const annotations = JSON.parse(jsonStr);

                    // 按照图像相对路径生成标注文件路径
                    imgRelativePath = path.relative(datasetDir, imagePath);
                    const [labelFileDir, labelFilePath] = this.getLabelPath(datasetDir, workDir, datasetRelativePath, imagePath);
                    // console.log(`labelFileDir: ${labelFileDir}`);
                    // console.log(`labelFilePath: ${labelFilePath}`);
                    // console.log('=================================\n\n');

                    // 创建目录并写入标注文件
                    await fs.mkdir(labelFileDir, { recursive: true });
                    await fs.writeFile(labelFilePath, JSON.stringify(annotations));

                    matchingPairs.push({
                        image_path: path.relative(datasetDir, imagePath),
                        label_path: path.relative(datasetDir, labelFilePath)
                    });

                    labeledCount++;
                } catch (e) {
                    console.error(`Error processing line in ${labelFile}:`, e);
                }
            }
        }

        // 保存匹配关系
        await this.saveMatchingPairs(datasetId, matchingPairs);

        return {
            code: 0,
            msg: {
                totalSamples: imageFiles.length,
                labeledSamples: labeledCount
            }
        };
    }

    getLabelPath(datasetDir, workDir, imageName, imageFilePath) {
        const imageRelativePath = path.relative(datasetDir, imageFilePath);
        const labelRelativeDir = path.dirname(imageRelativePath);
        const newLabelFileDir = path.join(workDir, labelRelativeDir);
        const jsonFileName = path.basename(imageName, path.extname(imageName)) + '.json';
        const imageAnnotationsPath = path.join(newLabelFileDir, jsonFileName);
        return [newLabelFileDir, imageAnnotationsPath];
    }
}

module.exports = PaddleDatasetWrapper;